package com.bruce.ai.alibaba.workflow.config;

import java.util.HashMap;

import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.agent.ReactAgent;
import com.bruce.ai.alibaba.workflow.node.AnalysisNode;
import com.bruce.ai.alibaba.workflow.node.ReportNode;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;


@Configuration
public class AgentWorkflowConfig {

    @Bean 
    public StateGraph dataAnalysisGraph(@Qualifier("analysisAgent") ReactAgent analysisAgent,@Qualifier("reportAgent")ReactAgent reportAgent) throws Exception{
        
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> strategies = new HashMap<>();
            strategies.put("input", new ReplaceStrategy());
            strategies.put("analysis_result", new ReplaceStrategy());
            strategies.put("final_result", new ReplaceStrategy());
            return strategies ;
        };

        StateGraph graph = new StateGraph(keyStrategyFactory);

        graph.addNode("analysis", node_async(new AnalysisNode()));
        graph.addNode("report", node_async(new ReportNode()));

        // 定义流程
        graph.addEdge(StateGraph.START, "analysis");
        graph.addEdge("analysis", "report");
        graph.addEdge("report", StateGraph.END);
        return graph;

    }

}
